import pygame
import numpy as np
import math
import random
from typing import List, Tuple, Dict, Any


class ArmEnv:
    # 常量定义
    SCREEN_WIDTH = 640
    SCREEN_HEIGHT = 640
    WORLD_WIDTH = 5.0  # 米
    WORLD_HEIGHT = 5.0  # 米
    SCALE = SCREEN_WIDTH / WORLD_WIDTH  # 像素/米

    # 颜色定义
    WHITE = (255, 255, 255)
    BLACK = (0, 0, 0)
    GRAY = (128, 128, 128)
    RED = (255, 0, 0)
    GREEN = (0, 255, 0)
    BLUE = (0, 120, 255)
    YELLOW = (255, 255, 0)
    PURPLE = (180, 0, 255)
    PINK = (255, 100, 180)
    ORANGE = (255, 150, 0)

    # 所有可用颜色列表（用于随机选择）
    ALL_COLORS = [RED, GREEN, BLUE, YELLOW, PURPLE, PINK, ORANGE]

    # 机械臂参数
    LINK_LENGTHS = [2.0, 1.6, 0.4]  # 三节长度
    LINK_ANGLE_ACC_LIMIT = 0.05

    def __init__(self):
        # 初始化pygame
        pygame.init()
        self.screen = pygame.display.set_mode((self.SCREEN_WIDTH, self.SCREEN_HEIGHT))
        pygame.display.set_caption("2D Robotic Arm Environment")
        self.clock = pygame.time.Clock()

        # 物理参数
        self.gravity = 9.8  # m/s² (y轴方向)
        self.dt = 0.1  # 时间步长
        self.elasticity = 0.8  # 碰撞弹性系数

        # 机械臂状态
        self.base_pos = np.array(
            [self.WORLD_WIDTH / 2, self.WORLD_HEIGHT / 2]
        )  # 基座位置
        self.joint_angles = [0.0, 0.0, 0.0]  # 三个关节角度 (弧度)
        self.gripper_state = False  # 夹爪状态 (False=松开, True=抓取)

        # 目标关节角度（用于插值）
        self.target_joint_angles = self.joint_angles.copy()

        # 环境物体
        self.tables = []
        self.boxes = []

        # 墙壁边界
        self.walls = [
            {"pos": np.array([0, 0]), "size": np.array([0.1, self.WORLD_HEIGHT])},  # 左墙
            {"pos": np.array([self.WORLD_WIDTH, 0]), "size": np.array([0.1, self.WORLD_HEIGHT])},  # 右墙
            {"pos": np.array([0, self.WORLD_HEIGHT]), "size": np.array([self.WORLD_WIDTH, 0.1])},  # 下墙（地面）
        ]

        # 重置环境
        self.reset()

    def reset(self):
        """重置环境到初始状态"""
        # 重置机械臂
        self.joint_angles = [0.0, 0.0, 0.0]
        self.target_joint_angles = self.joint_angles.copy()
        self.gripper_state = False

        # 清除现有物体
        self.tables = []
        self.boxes = []

        # 创建桌子 (3-5张)
        num_tables = random.randint(3, 5)
        for _ in range(num_tables):
            # 确保桌子不会重叠
            valid_position = False
            while not valid_position:
                x = random.uniform(0.5, self.WORLD_WIDTH - 0.5)
                y = random.uniform(0.5, self.WORLD_HEIGHT - 0.5)
                new_table = {
                    "pos": np.array([x, y]),
                    "size": np.array([0.5, 0.5]),
                    "color": random.choice(self.ALL_COLORS),  # 随机选择颜色
                }

                # 检查是否与现有桌子重叠
                overlap = False
                for table in self.tables:
                    if self._check_rect_overlap(
                        new_table["pos"], new_table["size"], table["pos"], table["size"]
                    ):
                        overlap = True
                        break

                if not overlap:
                    self.tables.append(new_table)
                    valid_position = True

        # 创建箱子 (4-6个)
        num_boxes = random.randint(4, 6)
        for _ in range(num_boxes):
            # 箱子可以在地面或桌子上
            if self.tables and random.random() > 0.3:  # 70%几率在桌子上
                table = random.choice(self.tables)
                x = table["pos"][0] + random.uniform(-0.15, 0.15)
                y = table["pos"][1] - table["size"][1] / 2 - 0.1
            else:  # 在地面上
                x = random.uniform(0.2, self.WORLD_WIDTH - 0.2)
                y = self.WORLD_HEIGHT - 0.1  # 地面高度

            self.boxes.append(
                {
                    "pos": np.array([x, y]),
                    "size": np.array([0.2, 0.2]),
                    "color": random.choice(self.ALL_COLORS),  # 随机选择颜色
                    "velocity": np.array([0.0, 0.0]),
                    "on_table": False,
                    "grabbed": False,
                }
            )

    def step(self, action: np.ndarray):
        """
        执行动作并更新环境
        action: 三个关节的角度变化量 [Δθ0, Δθ1, Δθ2, gripper_action]
                其中 gripper_action: 0=无动作, 1=抓取, 2=松开
        """
        # 更新目标关节角度
        for i in range(3):
            acc = np.clip(
                action[i], -self.LINK_ANGLE_ACC_LIMIT, self.LINK_ANGLE_ACC_LIMIT
            )
            self.joint_angles[i] += acc

        # 处理夹爪动作
        if action[3] == 1:  # 抓取
            self.gripper_state = True
        elif action[3] == 2:  # 松开
            self.gripper_state = False

        # 更新被抓取的箱子
        self._update_grabbed_box()

        # 更新物理模拟
        self._physics_update()

        # 获取观测值
        observation = self._get_observation()

        # 渲染环境
        self.render()

        return observation

    def _update_grabbed_box(self):
        """更新被抓取的箱子状态"""
        # 获取末端执行器位置
        end_effector_pos = self._forward_kinematics()
        
        # 计算有效抓取区域（箱子本身区域加上上方一个箱子高度的区域）
        for box in self.boxes:
            if not box["grabbed"]:
                # 计算抓取区域范围
                grab_area_left = box["pos"][0] - box["size"][0] / 2
                grab_area_right = box["pos"][0] + box["size"][0] / 2
                grab_area_bottom = box["pos"][1] + box["size"][1] / 2  # 箱子底部
                grab_area_top = box["pos"][1] - 1.5 * box["size"][1]   # 箱子顶部上方一个箱子高度
                
                # 检查末端执行器是否在抓取区域内
                in_grab_area = (
                    grab_area_left <= end_effector_pos[0] <= grab_area_right and
                    grab_area_top <= end_effector_pos[1] <= grab_area_bottom
                )
                
                # 如果夹爪是闭合的，检查是否有箱子被抓取
                if self.gripper_state and in_grab_area:
                    # 确保没有其他箱子已经被抓取
                    if not any(b["grabbed"] for b in self.boxes):
                        box["grabbed"] = True
                        box["velocity"] = np.array([0.0, 0.0])  # 重置速度
                        # 设置箱子位置为末端执行器下方一个箱子尺寸的位置
                        box["pos"] = end_effector_pos + np.array([0, box["size"][1] / 2])
        else:
            # 松开所有被抓取的箱子
            for box in self.boxes:
                if box["grabbed"]:
                    box["grabbed"] = False

    def _physics_update(self):
        """更新物理模拟（重力和碰撞）"""
        # 先应用重力并更新位置
        for box in self.boxes:
            if box["grabbed"]:
                # 被抓取的箱子跟随末端执行器 - 位于末端执行器下方一个箱子尺寸的位置
                end_effector_pos = self._forward_kinematics()
                box["pos"] = end_effector_pos + np.array([0, box["size"][1] / 2])
                box["velocity"] = np.array([0.0, 0.0])
                box["on_table"] = False
            elif not box["grabbed"]:  # 只有未被抓取的箱子受重力影响
                # 应用重力
                box["velocity"][1] += self.gravity * self.dt

                # 更新位置
                new_pos = box["pos"] + box["velocity"] * self.dt

                # 检查与地面的碰撞
                ground_level = self.WORLD_HEIGHT - box["size"][1] / 2
                if new_pos[1] > ground_level:
                    new_pos[1] = ground_level
                    box["velocity"][1] = 0.0
                    box["on_table"] = False

                # 检查与桌子的碰撞
                box["on_table"] = False
                for table in self.tables:
                    table_top = table["pos"][1] - table["size"][1] / 2
                    table_bottom = table["pos"][1] + table["size"][1] / 2

                    # 检查箱子是否在桌面上方
                    if (
                        new_pos[0] > table["pos"][0] - table["size"][0] / 2
                        and new_pos[0] < table["pos"][0] + table["size"][0] / 2
                        and new_pos[1] > table_top - box["size"][1] / 2
                        and new_pos[1] < table_bottom + box["size"][1] / 2
                    ):

                        # 碰撞响应
                        new_pos[1] = table_top - box["size"][1] / 2
                        box["velocity"][1] = 0.0
                        box["on_table"] = True
                        
                # 检查与墙壁的碰撞 - 仅适用于未被抓取的箱子
                for wall in self.walls:
                    # 检查箱子是否与墙壁碰撞
                    if self._check_rect_overlap(
                        new_pos, box["size"], wall["pos"], wall["size"]
                    ):
                        # 计算碰撞方向
                        dx = new_pos[0] - wall["pos"][0]
                        dy = new_pos[1] - wall["pos"][1]
                        
                        # 确定碰撞方向（水平或垂直）
                        if abs(dx) > abs(dy):  # 水平碰撞
                            # 反转水平速度
                            box["velocity"][0] = -box["velocity"][0] * self.elasticity
                            # 调整位置防止穿透
                            if dx > 0:  # 箱子在墙壁右侧
                                new_pos[0] = wall["pos"][0] + wall["size"][0]/2 + box["size"][0]/2
                            else:  # 箱子在墙壁左侧
                                new_pos[0] = wall["pos"][0] - wall["size"][0]/2 - box["size"][0]/2
                        else:  # 垂直碰撞
                            # 反转垂直速度
                            box["velocity"][1] = -box["velocity"][1] * self.elasticity
                            # 调整位置防止穿透
                            if dy > 0:  # 箱子在墙壁下方
                                new_pos[1] = wall["pos"][1] + wall["size"][1]/2 + box["size"][1]/2
                            else:  # 箱子在墙壁上方
                                new_pos[1] = wall["pos"][1] - wall["size"][1]/2 - box["size"][1]/2

                # 更新位置
                box["pos"] = new_pos
        
        # 处理箱子之间的碰撞（仅适用于未被抓取的箱子）
        self._handle_box_collisions()

    def _handle_box_collisions(self):
        """处理箱子之间的碰撞（仅处理未被抓取的箱子）"""
        # 检查所有箱子对
        for i in range(len(self.boxes)):
            for j in range(i + 1, len(self.boxes)):
                box1 = self.boxes[i]
                box2 = self.boxes[j]
                
                # 跳过被抓取的箱子
                if box1["grabbed"] or box2["grabbed"]:
                    continue
                
                # 检查两个箱子是否碰撞
                if self._check_rect_overlap(
                    box1["pos"], box1["size"], box2["pos"], box2["size"]
                ):
                    # 计算碰撞方向向量
                    collision_dir = box2["pos"] - box1["pos"]
                    distance = np.linalg.norm(collision_dir)
                    
                    if distance == 0:  # 避免除以零
                        collision_dir = np.array([1, 0])
                    else:
                        collision_dir = collision_dir / distance
                    
                    # 计算重叠深度
                    min_dist = (box1["size"][0] + box2["size"][0]) / 2
                    overlap = min_dist - distance
                    
                    # 分离箱子（防止重叠）
                    if not box1["grabbed"] and not box2["grabbed"]:
                        # 两个箱子都可以移动
                        box1["pos"] -= collision_dir * overlap / 2
                        box2["pos"] += collision_dir * overlap / 2
                        
                        # 计算相对速度
                        relative_velocity = box2["velocity"] - box1["velocity"]
                        velocity_along_normal = np.dot(relative_velocity, collision_dir)
                        
                        # 如果箱子正在分离，不处理碰撞
                        if velocity_along_normal > 0:
                            continue
                        
                        # 计算冲量
                        impulse = -(1 + self.elasticity) * velocity_along_normal
                        impulse /= 2  # 因为两个箱子质量相同
                        
                        # 应用冲量
                        impulse_vec = impulse * collision_dir
                        box1["velocity"] -= impulse_vec
                        box2["velocity"] += impulse_vec
                    
                    elif box1["grabbed"] and not box2["grabbed"]:
                        # 只有box1被抓取，只移动box2
                        box2["pos"] += collision_dir * overlap
                        
                        # 计算反射方向
                        dot_product = np.dot(box2["velocity"], collision_dir)
                        if dot_product < 0:  # 只处理朝向碰撞方向的运动
                            box2["velocity"] -= (1 + self.elasticity) * dot_product * collision_dir
                    
                    elif not box1["grabbed"] and box2["grabbed"]:
                        # 只有box2被抓取，只移动box1
                        box1["pos"] -= collision_dir * overlap
                        
                        # 计算反射方向
                        dot_product = np.dot(box1["velocity"], collision_dir)
                        if dot_product < 0:  # 只处理朝向碰撞方向的运动
                            box1["velocity"] -= (1 + self.elasticity) * dot_product * collision_dir

    def _forward_kinematics(self) -> np.ndarray:
        """计算机械臂末端执行器位置"""
        x = self.base_pos[0]
        y = self.base_pos[1]

        # 计算每个关节的位置
        angle = self.joint_angles[0]
        x += self.LINK_LENGTHS[0] * np.cos(angle)
        y += self.LINK_LENGTHS[0] * np.sin(angle)

        angle += self.joint_angles[1]
        x += self.LINK_LENGTHS[1] * np.cos(angle)
        y += self.LINK_LENGTHS[1] * np.sin(angle)

        angle += self.joint_angles[2]
        x += self.LINK_LENGTHS[2] * np.cos(angle)
        y += self.LINK_LENGTHS[2] * np.sin(angle)

        return np.array([x, y])

    def _get_observation(self) -> Dict[str, Any]:
        """获取环境观测值"""
        # 1. 获取RGB图像
        self.render()
        rgb_array = pygame.surfarray.array3d(self.screen)
        # 转换维度顺序为 (C, H, W)
        rgb_array = np.transpose(rgb_array, (2, 0, 1))

        # 2. 获取物体状态信息
        objects = []

        # 添加桌子信息
        for table in self.tables:
            objects.append(
                [
                    "table",
                    table["pos"][0],
                    table["pos"][1],
                    0.0,  # sin(rotation)
                    1.0,  # cos(rotation)
                    table["color"],
                ]
            )

        # 添加箱子信息
        for box in self.boxes:
            objects.append(
                [
                    "box",
                    box["pos"][0],
                    box["pos"][1],
                    0.0,  # sin(rotation)
                    1.0,  # cos(rotation)
                    box["color"],
                ]
            )

        # 3. 获取机械臂状态
        arm_state = {
            "base_pos": self.base_pos.tolist(),
            "joint_angles": self.joint_angles.copy(),
            "end_effector_pos": self._forward_kinematics().tolist(),
            "gripper_state": self.gripper_state,
        }

        return {"rgb_image": rgb_array, "objects": objects, "arm_state": arm_state}

    def render(self):
        """渲染环境"""
        self.screen.fill(self.WHITE)

        # 绘制墙壁
        pygame.draw.rect(self.screen, self.BLACK, (0, 0, self.SCREEN_WIDTH, 5))
        pygame.draw.rect(self.screen, self.BLACK, (0, 0, 5, self.SCREEN_HEIGHT))
        pygame.draw.rect(
            self.screen, self.BLACK, (self.SCREEN_WIDTH - 5, 0, 5, self.SCREEN_HEIGHT)
        )
        pygame.draw.rect(
            self.screen, self.BLACK, (0, self.SCREEN_HEIGHT - 5, self.SCREEN_WIDTH, 5)
        )

        # 绘制桌子
        for table in self.tables:
            x, y = self._world_to_screen(table["pos"])
            w, h = table["size"] * self.SCALE
            pygame.draw.rect(self.screen, table["color"], (x - w / 2, y - h / 2, w, h))

        # 绘制箱子
        for box in self.boxes:
            x, y = self._world_to_screen(box["pos"])
            w, h = box["size"] * self.SCALE
            pygame.draw.rect(self.screen, box["color"], (x - w / 2, y - h / 2, w, h))

        # 绘制机械臂
        self._draw_arm()

        pygame.display.flip()
        self.clock.tick(30)  # 30 FPS

    def _draw_arm(self):
        """绘制机械臂"""
        # 基座位置
        base_screen = self._world_to_screen(self.base_pos)

        # 计算各关节位置
        positions = [self.base_pos]
        angle = self.joint_angles[0]
        pos1 = self.base_pos + np.array(
            [self.LINK_LENGTHS[0] * np.cos(angle), self.LINK_LENGTHS[0] * np.sin(angle)]
        )
        positions.append(pos1)

        angle += self.joint_angles[1]
        pos2 = pos1 + np.array(
            [self.LINK_LENGTHS[1] * np.cos(angle), self.LINK_LENGTHS[1] * np.sin(angle)]
        )
        positions.append(pos2)

        angle += self.joint_angles[2]
        end_effector = pos2 + np.array(
            [self.LINK_LENGTHS[2] * np.cos(angle), self.LINK_LENGTHS[2] * np.sin(angle)]
        )
        positions.append(end_effector)

        # 转换为屏幕坐标
        screen_positions = [self._world_to_screen(p) for p in positions]

        # 绘制连杆
        pygame.draw.line(
            self.screen, self.BLACK, screen_positions[0], screen_positions[1], 8
        )
        pygame.draw.line(
            self.screen, self.GRAY, screen_positions[1], screen_positions[2], 6
        )
        pygame.draw.line(
            self.screen, self.GRAY, screen_positions[2], screen_positions[3], 4
        )

        # 绘制关节
        for i, pos in enumerate(screen_positions[:3]):
            pygame.draw.circle(self.screen, self.BLACK, pos, 10)

        # 绘制末端执行器
        if self.gripper_state:  # 抓取状态
            pygame.draw.circle(self.screen, self.BLACK, screen_positions[3], 8)
            # 绘制被抓取的箱子（在末端执行器下方）
            if any(box["grabbed"] for box in self.boxes):
                # 找到被抓取的箱子
                grabbed_box = next(box for box in self.boxes if box["grabbed"])
                # 绘制箱子在末端执行器下方
                box_pos = self._world_to_screen(grabbed_box["pos"])
                w, h = grabbed_box["size"] * self.SCALE
                pygame.draw.rect(
                    self.screen, 
                    grabbed_box["color"], 
                    (box_pos[0] - w / 2, box_pos[1] - h / 2, w, h)
                )
        else:  # 松开状态
            pygame.draw.circle(self.screen, self.BLACK, screen_positions[3], 8, 2)

    def _world_to_screen(self, world_pos: np.ndarray) -> Tuple[int, int]:
        """世界坐标转换为屏幕坐标"""
        x = int(world_pos[0] * self.SCALE)
        y = int(world_pos[1] * self.SCALE)
        return (x, y)

    def _wrap_angle(self, angle: float) -> float:
        """将角度限制在[-π, π]范围内"""
        while angle > np.pi:
            angle -= 2 * np.pi
        while angle < -np.pi:
            angle += 2 * np.pi
        return angle

    def _check_rect_overlap(self, pos1, size1, pos2, size2):
        """检查两个矩形是否重叠"""
        return (abs(pos1[0] - pos2[0]) < (size1[0] / 2 + size2[0] / 2)) and (
            abs(pos1[1] - pos2[1]) < (size1[1] / 2 + size2[1] / 2)
        )

    def close(self):
        """关闭环境"""
        pygame.quit()


# 测试环境
if __name__ == "__main__":
    env = ArmEnv()
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # 随机动作测试
        action = np.array(
            [
                random.uniform(-0.1, 0.1),  # Δθ0
                random.uniform(-0.1, 0.1),  # Δθ1
                random.uniform(-0.1, 0.1),  # Δθ2
                random.choice([0, 1, 2]),  # 夹爪动作
            ]
        )

        env.step(action)

    env.close()